package com.kexio.websocket.interceptor;

import cn.dev33.satoken.stp.StpUtil;
import cn.hutool.core.util.StrUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * WebSocket 认证拦截器
 * 
 * 在 WebSocket 握手阶段进行用户认证和权限校验。
 * 集成 Sa-Token 进行 token 验证。
 * 
 * @author Kexio Team
 */
@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {

    private static final Logger log = LoggerFactory.getLogger(WebSocketAuthInterceptor.class);

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        
        String uri = request.getURI().toString();
        log.debug("WebSocket 握手请求: {}", uri);
        
        try {
            // 从请求参数中获取认证信息
            String query = request.getURI().getQuery();
            Map<String, String> params = parseQueryString(query);
            
            String token = params.get("token");
            
            // 验证必要参数
            if (StrUtil.isBlank(token)) {
                log.warn("WebSocket 握手失败: 缺少 token 参数");
                return false;
            }
            
            // 使用 Sa-Token 验证 token 并获取用户信息
            String userId;
            String tenantId = null;
            
            try {
                // 从 token 中获取登录ID
                Object loginId = StpUtil.getLoginIdByToken(token);
                if (loginId == null) {
                    log.warn("WebSocket 握手失败: Token 无效");
                    return false;
                }
                
                // 解析 loginId（格式可能是 "userId" 或 "tenantId:userId"）
                String loginIdStr = loginId.toString();
                if (loginIdStr.contains(":")) {
                    String[] parts = loginIdStr.split(":", 2);
                    tenantId = parts[0];
                    userId = parts[1];
                } else {
                    userId = loginIdStr;
                }
                
                log.debug("Token 验证成功: userId={}, tenantId={}", userId, tenantId);
                
            } catch (Exception e) {
                log.warn("WebSocket 握手失败: Token 验证异常, error={}", e.getMessage());
                return false;
            }
            
            // 将用户信息存储到会话属性中
            attributes.put("userId", userId);
            attributes.put("tenantId", tenantId != null ? tenantId : "default");
            attributes.put("token", token);
            attributes.put("connectTime", System.currentTimeMillis());
            
            log.info("WebSocket 握手成功: userId={}, tenantId={}", userId, tenantId);
            return true;
            
        } catch (Exception e) {
            log.error("WebSocket 握手异常: {}", e.getMessage(), e);
            return false;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        if (exception != null) {
            log.error("WebSocket 握手后异常: {}", exception.getMessage(), exception);
        } else {
            log.debug("WebSocket 握手完成");
        }
    }


    /**
     * 解析查询字符串参数
     * 
     * @param query 查询字符串
     * @return 参数映射
     */
    private Map<String, String> parseQueryString(String query) {
        Map<String, String> params = new java.util.HashMap<>();
        
        if (StrUtil.isNotBlank(query)) {
            String[] pairs = query.split("&");
            for (String pair : pairs) {
                String[] keyValue = pair.split("=", 2);
                if (keyValue.length == 2) {
                    params.put(keyValue[0], keyValue[1]);
                }
            }
        }
        
        return params;
    }
}
